"""
Transformation Rule Registry.

This module provides a registry system for managing and applying transformation rules.
"""

from typing import Any

from orchestrator.transforms.rules import TransformationPhase, TransformationRule


class TransformationRegistry:
    """
    Registry for transformation rules.

    The registry maintains a collection of transformation rules and provides
    methods to apply them in the correct order.
    """

    def __init__(self) -> None:
        """Initialize the transformation registry"""
        self._rules: dict[TransformationPhase, list[TransformationRule]] = {
            phase: [] for phase in TransformationPhase
        }
        self._rules_by_name: dict[str, TransformationRule] = {}

    def register(self, rule: TransformationRule) -> None:
        """
        Register a transformation rule.

        Args:
            rule: The transformation rule to register
        """
        phase = rule.get_phase()
        self._rules[phase].append(rule)
        self._rules_by_name[rule.metadata.name] = rule

        # Sort rules by priority (lower priority executes first)
        self._rules[phase].sort(key=lambda r: r.get_priority())

    def unregister(self, rule_name: str) -> None:
        """
        Unregister a transformation rule by name.

        Args:
            rule_name: Name of the rule to unregister
        """
        if rule_name not in self._rules_by_name:
            return

        rule = self._rules_by_name[rule_name]
        phase = rule.get_phase()
        self._rules[phase] = [r for r in self._rules[phase] if r.metadata.name != rule_name]
        del self._rules_by_name[rule_name]

    def get_rule(self, rule_name: str) -> TransformationRule | None:
        """
        Get a rule by name.

        Args:
            rule_name: Name of the rule

        Returns:
            The rule or None if not found
        """
        return self._rules_by_name.get(rule_name)

    def get_rules_for_phase(self, phase: TransformationPhase) -> list[TransformationRule]:
        """
        Get all rules for a specific phase.

        Args:
            phase: The transformation phase

        Returns:
            List of rules for the phase, sorted by priority
        """
        return [r for r in self._rules[phase] if r.is_enabled()]

    def apply_phase(self, node: Any, phase: TransformationPhase) -> Any:
        """
        Apply all rules for a specific phase to a node.

        Args:
            node: The IR node to transform
            phase: The transformation phase

        Returns:
            Transformed node
        """
        current_node = node
        rules = self.get_rules_for_phase(phase)

        for rule in rules:
            if rule.can_apply(current_node):
                current_node = rule.apply(current_node)

        return current_node

    def apply_all_phases(self, node: Any) -> Any:
        """
        Apply all transformation phases to a node.

        Args:
            node: The IR node to transform

        Returns:
            Transformed node
        """
        current_node = node

        # Apply phases in order
        for phase in TransformationPhase:
            current_node = self.apply_phase(current_node, phase)

        return current_node

    def list_rules(self) -> list[str]:
        """
        List all registered rule names.

        Returns:
            List of rule names
        """
        return list(self._rules_by_name.keys())

    def get_rule_count(self) -> int:
        """
        Get the total number of registered rules.

        Returns:
            Number of rules
        """
        return len(self._rules_by_name)

    def clear(self) -> None:
        """Clear all registered rules"""
        self._rules = {phase: [] for phase in TransformationPhase}
        self._rules_by_name = {}


# Global registry instance
_global_registry: TransformationRegistry | None = None


def get_global_registry() -> TransformationRegistry:
    """
    Get the global transformation registry.

    Returns:
        The global registry instance
    """
    global _global_registry
    if _global_registry is None:
        _global_registry = TransformationRegistry()
    return _global_registry


def reset_global_registry() -> None:
    """Reset the global transformation registry"""
    global _global_registry
    _global_registry = None
